import logging
import argparse
from local_model import LocalModel
from prompt_matrix import PromptMatrix
from multimodal_processor import MultimodalProcessor
from utils import save_result, compare_results, timer_decorator, create_example

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def run_single_generation(content, parameters=None, model_name="runwayml/stable-diffusion-v1-5"):
    """运行单次生成任务"""
    if parameters is None:
        parameters = {
            "tone": "professional ",
            "detail_level": "middle ",
            "format": "true ",
            "perspective": "objective "
        }

    # 创建提示词矩阵实例
    prompt_matrix = PromptMatrix()

    # 生成提示词
    prompt = prompt_matrix.generate_prompt(
        base_content=content,
        parameters=parameters,
        template_type="with_perspective"
    )

    if not prompt:
        logger.error("无法生成提示词，任务终止")
        return None

    logger.info(f"生成的提示词: {prompt}")

    # 使用本地模型生成图像
    model = LocalModel(model_name=model_name)
    image_path = model.generate_image(prompt)

    if not image_path:
        logger.error("图像生成失败，任务终止")
        return None

    # 整理结果
    result = {
        "parameters": parameters,
        "prompt": prompt,
        "output": f"图像已生成: {image_path}",
        "image_path": image_path
    }

    return result


@timer_decorator
def run_matrix_experiment(content, categories=None, limit=10, model_name="runwayml/stable-diffusion-v1-5"):
    """运行提示词矩阵实验"""
    # 创建提示词矩阵实例
    prompt_matrix = PromptMatrix()

    # 生成提示词组合
    combinations = prompt_matrix.generate_matrix_combinations(include_categories=categories)

    # 如果组合太多，限制数量
    if len(combinations) > limit:
        logger.info(f"提示词组合太多 ({len(combinations)}), 将限制为 {limit} 个")
        combinations = combinations[:limit]

    # 运行实验
    results = []
    model = LocalModel(model_name=model_name)

    for i, params in enumerate(combinations, 1):
        logger.info(f"运行实验 {i}/{len(combinations)}，参数: {params}")

        # 生成提示词
        prompt = prompt_matrix.generate_prompt(
            base_content=content,
            parameters=params,
            template_type="with_perspective"
        )

        if not prompt:
            logger.warning(f"跳过无效参数组合: {params}")
            continue

        # 生成图像
        image_path = model.generate_image(prompt)

        if image_path:
            results.append({
                "parameters": params,
                "prompt": prompt,
                "output": f"图像已生成: {image_path}",
                "image_path": image_path
            })
        else:
            logger.warning(f"图像生成失败，参数组合: {params}")

    # 保存结果
    prompt_matrix.save_matrix_results(results)

    # 分析结果
    if len(results) >= 2:
        analysis = compare_results(results)
        save_result(analysis, prefix="matrix_analysis")
        logger.info("提示词矩阵实验分析已完成")

    return results


def main():
    """主函数"""
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="多模态内容生成系统")
    parser.add_argument("--mode", choices=["single", "matrix"], default="single",
                        help="运行模式: single(单次生成) 或 matrix(提示词矩阵实验)")
    parser.add_argument("--input", help="输入内容或文件路径")
    parser.add_argument("--limit", type=int, default=10,
                        help="提示词矩阵实验的最大组合数量")
    parser.add_argument("--create-example", help="创建示例数据")
    parser.add_argument("--model", default="runwayml/stable-diffusion-v1-5",
                        help="HuggingFace模型名称")

    args = parser.parse_args()

    # 处理示例创建
    if args.create_example:
        example_data = {
            "text": "一只可爱的猫",
            "images": []
        }
        create_example(args.create_example, example_data)
        return

    # 处理输入
    processor = MultimodalProcessor()

    if args.input:
        content = processor.process_input(args.input)
    else:
        # 使用默认内容
        logger.info("未指定输入，使用默认示例内容")
        content = "一只可爱的猫"

    if not content:
        logger.error("无法处理输入内容，程序退出")
        return

    # 根据模式运行
    if args.mode == "single":
        logger.info("运行单次生成任务...")
        result = run_single_generation(content, model_name=args.model)

        if result:
            save_result(result, prefix="single_generation")
            logger.info(f"生成完成，图像保存至: {result['image_path']}")

    elif args.mode == "matrix":
        logger.info("运行提示词矩阵实验...")
        results = run_matrix_experiment(content, limit=args.limit, model_name=args.model)
        logger.info(f"提示词矩阵实验完成，共生成 {len(results)} 个结果")


if __name__ == "__main__":
    main()